#coding=utf-8
import numpy as np
import os,sys
from tqdm import tqdm
sys.path.append('../../')
import clean_operations_py3 as cleop
import cv2
import shutil
import xml.etree.ElementTree as ET
from xml.dom.minidom import Document

#将一个文件夹中的所有xml转成对应的txt
def convert_xml_to_ssd_label():
    xml_root='./annotations_xmls'

    xml_list=cleop.get_target_files(xml_root,rules=['.xml'])
    length=len(xml_list)
    print(length)

    for i in tqdm(range(length)):
        xml_path=xml_list[i]
        txt_path =xml_path.replace('.xml','.txt').replace('annotations_xmls','labels')
        print(xml_path)
        img_path=xml_path.replace('.xml','.jpg').replace('annotations_xmls','images')
        if not os.path.exists(img_path):
            continue

        img=cv2.imread(img_path)
        height0,width0,_=img.shape
        # print(width,height)

        root=ET.parse(xml_path).getroot()

        size=root.find('size')
        width=int(size[0].text)
        height=int(size[1].text)
        # print(width,height)
        #assert(height0==height and width0==width)

        txt_direct=cleop.get_file_direct(txt_path)
        cleop.make_direct_exist(txt_direct)
        fw=open(txt_path,'w')

        for child in root.findall('object'):
            name = child.find('name').text
            #print(name == 'head')
            sub=child.find('bndbox')
            xmin=float(sub[0].text)
            ymin=float(sub[1].text)
            xmax=float(sub[2].text)
            ymax=float(sub[3].text)
            # print(xmin,ymin,xmax,ymax)
            xmin=int(xmin)
            ymin=int(ymin)
            xmax=int(xmax)
            ymax=int(ymax)

            xmin=xmin if xmin>=0 else 0
            ymin=ymin if ymin>=0 else 0
            xmax=xmax if xmax <width else width-1
            ymax=ymax if ymax <height else height-1

            # cx=round((xmin+xmax)/2/width,6)
            # cy=round((ymin+ymax)/2/height,6)
            # cw=round((xmax-xmin)/width,6)
            # ch=round((ymax-ymin)/height,6)
            # # type=0
            # print(cx,cy,cw,ch)
            
            if name == 'fire':
                fw.write('%d %d %d %d %d\n' %(1,xmin,ymin,xmax,ymax))
            # elif name == 'smoke':
            #     fw.write('%d %d %d %d %d\n' %(2,xmin,ymin,xmax,ymax))
            else:
                print('Error')
            #-------------------------------------
            # if name == 'phone':
            #     fw.write('%d %f %f %f %f\n' %(0,cx,cy,cw,ch))
            # elif name == 'smoke':
            #     fw.write('%d %f %f %f %f\n' %(1,cx,cy,cw,ch))
            # else:
            #     print('Error')
        fw.close()


def convert_xml_to_yolo_label():
    xml_root='./Annotations'

    xml_list=cleop.get_target_files(xml_root,rules=['.xml'])
    length=len(xml_list)
    print(length)

    for i in tqdm(range(length)):
        xml_path=xml_list[i]
        txt_path =xml_path.replace('.xml','.txt').replace('Annotations','Labels')
        print(xml_path)
        img_path=xml_path.replace('.xml','.jpg').replace('Annotations','Images')
        if not os.path.exists(img_path):
            continue

        img=cv2.imread(img_path)
        height0,width0,_=img.shape
        # print(width,height)

        root=ET.parse(xml_path).getroot()

        size=root.find('size')
        width=int(size[0].text)
        height=int(size[1].text)
        # print(width,height)
        #assert(height0==height and width0==width)

        txt_direct=cleop.get_file_direct(txt_path)
        cleop.make_direct_exist(txt_direct)
        fw=open(txt_path,'w')

        for child in root.findall('object'):
            name = child.find('name').text
            #print(name == 'head')
            sub=child.find('bndbox')
            xmin=float(sub[0].text)
            ymin=float(sub[1].text)
            xmax=float(sub[2].text)
            ymax=float(sub[3].text)
            # print(xmin,ymin,xmax,ymax)
            xmin=int(xmin)
            ymin=int(ymin)
            xmax=int(xmax)
            ymax=int(ymax)

            xmin=xmin if xmin>=0 else 0
            ymin=ymin if ymin>=0 else 0
            xmax=xmax if xmax <width else width-1
            ymax=ymax if ymax <height else height-1

            cx=round((xmin+xmax)/2/width,6)
            cy=round((ymin+ymax)/2/height,6)
            cw=round((xmax-xmin)/width,6)
            ch=round((ymax-ymin)/height,6)
            # # type=0
            # print(cx,cy,cw,ch)
            
            #-------------------------------------
            if name == 'fire':
                fw.write('%d %f %f %f %f\n' %(0,cx,cy,cw,ch))
            else:
                print('Error')
        fw.close()

if __name__ == '__main__':
    convert_xml_to_ssd_label()
    #convert_xml_to_yolo_label()